"""
铁路货运量的时序建模
"""
import os
import datetime
import pandas as pd
from pylab import mpl
import matplotlib.pyplot as plt
from algorithm.line_model import line_model
from algorithm.stl_model import stl_model
from algorithm.holt_winters_damp_model import holt_winter_damped
from algorithm.arma_arima_model import arma_arima_model
from algorithm.lstm_model import lstm_model
# from algorithm.lstm_model_2 import lstm_model
from algorithm.gru_model import gru_model


# 指定默认字体
mpl.rcParams['font.sans-serif'] = ['SimHei']
# 解决保存图像是负号'-'显示为方块的问题
mpl.rcParams['axes.unicode_minus'] = False


def load_data():
    """
    数据加载，并生成训练集和测试集
    :return:
    """
    file = '../../data/运输量.xlsx'
    df = pd.read_excel(file, sheet_name=1)
    df['month'] = df['month'].apply(
        lambda x: datetime.datetime.strptime(x, '%Y年%m月'))
    df = df.sort_values(by='month')
    # df.index = pd.date_range(start='2015', periods=len(df), freq='M')
    df = df.set_index('month')
    df = df.sort_index()
    return df['铁路货运量当期值(万吨)']


def run():
    """
    主入口
    :return:
    """
    pic_save_path = r'C:\Users\xiwsu11\Desktop\时序分享\Python'
    # train_data, test_data = load_data()
    ts_data = load_data()

    # 数据探索
    """
    ts_data.plot()
    plt.xlabel('日期')
    plt.ylabel('铁路货运量(万吨)')
    plt.title('2015-2019年全年铁路货运量趋势图')
    pic_name = r'数据探索.png'
    plt.savefig(os.path.join(pic_save_path, pic_name))
    # plt.show()
    """

    # 建模
    # 简单线性回归
    # line_model(data=ts_data, pic_save_path=pic_save_path)

    # stl
    # stl_model(data=ts_data, pic_save_path=pic_save_path)

    # holt-winter季节性阻尼方法
    # holt_winter_damped(data=ts_data, pic_save_path=pic_save_path)

    #  ARIMA模型
    # arma_arima_model(data=ts_data, pic_save_path=pic_save_path)

    # LSTM模型
    # lstm_model(data=ts_data)

    # GRU
    gru_model(data=ts_data)


if __name__ == '__main__':
    run()
